#refer to the Example 4.5 in Chapter 4 "Forward Kinematics" in book 《Modern Robotics》

import numpy as np
import math

def rotMat(w,q):
    eMat=np.eye(3)+math.sin(q)*ssMat(w)+(1-math.cos(q))*ssMat(w)@ssMat(w)
    return eMat

def ssMat(w):
    a=np.zeros((3,3))
    a[0,1]=-w[2]
    a[1,0]= w[2]
    a[0,2]= w[1]
    a[2,0]=-w[1]
    a[1,2]=-w[0]
    a[2,1]= w[0]
    return a

def tranVec(w,v,q):
    tran_vec = (np.eye(3)*q+(1-math.cos(q))*ssMat(w)+(q-math.sin(q))*ssMat(w)@ssMat(w))@v;
    a=(q-math.sin(q))*ssMat(w)@ssMat(w)
    return tran_vec

def TMat(w,v,q):
    tMat=np.zeros((4,4))
    tMat[0:3,0:3]=rotMat(w,q)
    tMat[0:3,3:4]=tranVec(w,v,q)
    tMat[3,3]=1
    return tMat
    
def main():
    # modeling with POE algorithm 
    v2 = np.transpose(np.array([[-0.089,0,0]]))
    v5 = np.transpose(np.array([[-0.109,0.425+0.392,0]]))

    w2 = np.transpose(np.array([[0,1, 0]]))
    w5 = np.transpose(np.array([[0,0,-1]]))
    
    # specify the M matrix
    M  = np.array([[-1, 0, 0, 0.425+0.392],
                   [ 0, 0, 1, 0.109+0.082],
                   [ 0, 1, 0, 0.089-0.095],
                   [ 0, 0, 0, 1]])
    
    # given joint space input
    q2=-math.pi/2;
    q5= math.pi/2;
    
    # calculate matrix exponentials
    e2=TMat(w2,v2,q2)
    e5=TMat(w5,v5,q5)
    
    
    print("e2\n",e2)
    print("e5\n",e5)
    
    # get final homogeneous transformation matrix
    t=e2@e5@M
    print("t\n",t)

if __name__ == "__main__":
    main()